from ..builder import DETECTORS
from .two_stage import TwoStageDetector


@DETECTORS.register_module()
class WeakRCNN_FT(TwoStageDetector):
    """Implementation of `Fast R-CNN <https://arxiv.org/abs/1504.08083>`_"""

    def __init__(self,
                 backbone,
                 neck,
                 roi_head,
                 train_cfg,
                 test_cfg,
                 pretrained=None):
        super(WeakRCNN_FT, self).__init__(
            backbone=backbone,
            neck=neck,
            roi_head=roi_head,
            train_cfg=train_cfg,
            test_cfg=test_cfg,
            pretrained=pretrained)
    def extract_feat(self, img,ft):
        """Directly extract features from the backbone+neck."""
        x = self.backbone(img,ft)
#         if self.with_neck:
#             x = self.neck(x)
        return x
    
    def forward_train(self,
                      img,
                      img_metas,
                      gt_labels,
                      proposals=None,
                      **kwargs):

        try:
            ft_mean=[104., 117., 124.]
            ft = kwargs['ft']
            ft = ft.permute(0,3,1,2).float()
            import torch.nn as nn
            ft = nn.functional.interpolate(ft, (img.size(2), img.size(3)), 
mode='bilinear', align_corners=True)
#             ft[:,0,:,:] = ft[:,0,:,:]-ft_mean[0]
#             ft[:,1,:,:] = ft[:,1,:,:]-ft_mean[1]
#             ft[:,2,:,:] = ft[:,2,:,:]-ft_mean[2]

            #插值到图片大小

            import torch
            #concat之后是6通道
#             img_ft = torch.cat((img,ft),dim=1)

            
#             print(ft[0][0])
#             print(ft[0][1])
#             print(ft[0][2])            
#             print('ft',torch.max(ft))
#             print('ft',torch.min(ft))
#             print('img',torch.max(img))
#             print('img',torch.min(img))
            x = self.extract_feat(img,ft) 
            
        except:
            x = self.extract_feat(img)        
        losses = dict()

        proposal_list = proposals
        try:
            roi_losses = self.roi_head.forward_train(x, img_metas, proposal_list,
                                                     gt_labels, **kwargs)
        except:
            roi_losses = self.roi_head.forward_train(x, img_metas, proposal_list,
                                                     gt_labels)            
        losses.update(roi_losses)
        return losses


    def forward_test(self, imgs, img_metas, proposals, **kwargs):
        for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]:
            if not isinstance(var, list):
                raise TypeError(f'{name} must be a list, but got {type(var)}')
#         print(kwargs['ft'][0].shape)
        num_augs = len(imgs)
        if num_augs != len(img_metas):
            raise ValueError(f'num of augmentations ({len(imgs)}) '
                             f'!= num of image meta ({len(img_metas)})')

        if num_augs == 1:
            return self.simple_test(imgs[0], img_metas[0],kwargs['ft'][0], proposals[0],
                                    )
        else:
            assert imgs[0].size(0) == 1, 'aug test does not support ' \
                                         'inference with batch size ' \
                                         f'{imgs[0].size(0)}'
            return self.aug_test(imgs, img_metas, proposals, **kwargs)

    def simple_test(self, img, img_metas,ft, proposals=None, rescale=False):
        """Test without augmentation."""

        assert self.with_bbox, 'Bbox head must be implemented.'
        #
        ft_mean=[104., 117., 124.]
        ft = ft.permute(0,3,1,2).float()
        import torch.nn as nn
        ft = nn.functional.interpolate(ft, (img.size(2), img.size(3)), 
mode='bilinear', align_corners=True)
#         ft[:,0,:,:] = ft[:,0,:,:]-ft_mean[0]
#         ft[:,1,:,:] = ft[:,1,:,:]-ft_mean[1]
#         ft[:,2,:,:] = ft[:,2,:,:]-ft_mean[2]


        import torch
        #concat之后是6通道
#         img_ft = torch.cat((img,ft),dim=1)
        x = self.extract_feat(img,ft)
        if proposals is None:
            proposal_list = self.rpn_head.simple_test_rpn(x, img_metas)
        else:
            proposal_list = proposals

        return self.roi_head.simple_test(
            x, proposal_list, img_metas, rescale=rescale)
    
    def aug_test(self, imgs, img_metas, proposal_list, rescale=False):
        """Test with augmentations.

        If rescale is False, then returned bboxes and masks will fit the scale
        of imgs[0].
        """
        x = self.extract_feats(imgs)
        return self.roi_head.aug_test(
            x, proposal_list, img_metas, rescale=rescale)
